import pandas as pd
import numpy as np
import json
import sys
from sklearn.ensemble import RandomForestRegressor
from sklearn.metrics import mean_squared_error, mean_absolute_error, r2_score
import warnings
warnings.filterwarnings('ignore')

def convert_to_serializable(obj):
    """将numpy类型转换为Python原生类型以支持JSON序列化"""
    if isinstance(obj, np.integer):
        return int(obj)
    elif isinstance(obj, np.floating):
        return float(obj)
    elif isinstance(obj, np.ndarray):
        return obj.tolist()
    elif isinstance(obj, dict):
        return {key: convert_to_serializable(value) for key, value in obj.items()}
    elif isinstance(obj, list):
        return [convert_to_serializable(item) for item in obj]
    else:
        return obj

def random_forest_imputation(csv_file_path):
    """使用随机森林进行数据填充"""
    try:
        # 读取CSV数据
        df = pd.read_csv(csv_file_path)

        print(f"数据形状: {df.shape}")
        print(f"缺失值统计:")
        missing_info = df.isnull().sum()
        print(missing_info[missing_info > 0])

        # 记录原始数据统计
        original_stats = {
            'total_samples': int(len(df)),
            'total_features': int(df.shape[1]),
            'missing_values_count': int(df.isnull().sum().sum()),
            'missing_rate': float(df.isnull().sum().sum() / (df.shape[0] * df.shape[1]))
        }

        # 创建数据副本进行填充
        df_imputed = df.copy()
        imputation_metrics = {}

        # 对每个有缺失值的列进行随机森林填充
        for column in df.columns:
            if df[column].isnull().sum() > 0:
                print(f"正在填充列: {column}")

                # 分离有值和缺失值的数据
                train_data = df[df[column].notna()]
                missing_data = df[df[column].isna()]

                if len(train_data) == 0:
                    # 如果整列都是缺失值，用0填充
                    df_imputed[column] = df_imputed[column].fillna(0)
                    continue

                # 准备特征（除了当前列的其他列）
                feature_columns = [col for col in df.columns if col != column]
                X_train = train_data[feature_columns]
                y_train = train_data[column]

                # 处理训练数据中的缺失值（用均值填充）
                X_train = X_train.fillna(X_train.mean())

                # 如果缺失值行数过少，跳过模型训练
                if len(missing_data) == 0:
                    continue

                X_missing = missing_data[feature_columns]
                X_missing = X_missing.fillna(X_train.mean())  # 用训练集的均值填充

                # 训练随机森林模型
                try:
                    rf_model = RandomForestRegressor(
                        n_estimators=50,  # 减少树的数量以提高速度
                        random_state=42,
                        n_jobs=1,  # 使用单线程避免内存问题
                        max_depth=10,
                        min_samples_split=5,
                        min_samples_leaf=2
                    )

                    rf_model.fit(X_train, y_train)

                    # 预测缺失值
                    predictions = rf_model.predict(X_missing)

                    # 评估模型性能（使用交叉验证的方式）
                    if len(train_data) > 50:  # 只有足够数据时才进行评估
                        # 使用80%的训练数据训练，20%测试
                        split_idx = int(0.8 * len(train_data))
                        if split_idx > 0 and split_idx < len(train_data):
                            X_val = X_train.iloc[split_idx:]
                            y_val = y_train.iloc[split_idx:]
                            X_train_subset = X_train.iloc[:split_idx]
                            y_train_subset = y_train.iloc[:split_idx]

                            rf_model_val = RandomForestRegressor(
                                n_estimators=50,
                                random_state=42,
                                n_jobs=1,
                                max_depth=10
                            )
                            rf_model_val.fit(X_train_subset, y_train_subset)
                            val_predictions = rf_model_val.predict(X_val)

                            imputation_metrics[column] = {
                                'mse': float(mean_squared_error(y_val, val_predictions)),
                                'mae': float(mean_absolute_error(y_val, val_predictions)),
                                'r2': float(r2_score(y_val, val_predictions)),
                                'imputed_count': int(len(predictions))
                            }

                    # 填充缺失值
                    df_imputed.loc[df[column].isna(), column] = predictions

                    print(f"列 {column} 填充完成，填充了 {len(predictions)} 个值")

                except Exception as e:
                    print(f"列 {column} 随机森林填充失败，使用均值填充: {str(e)}")
                    # 如果随机森林失败，使用均值填充作为备选
                    mean_value = train_data[column].mean()
                    df_imputed.loc[df[column].isna(), column] = mean_value

        # 生成输出文件路径
        import os
        output_dir = "temp/output"
        os.makedirs(output_dir, exist_ok=True)

        output_csv_path = os.path.join(output_dir, f"rf_imputed_data.csv") # _{pd.Timestamp.now().strftime('%Y%m%d_%H%M%S')}

        # 保存填充后的数据
        df_imputed.to_csv(output_csv_path, index=False)

        # 计算填充后统计信息
        final_stats = {
            'total_samples': int(len(df_imputed)),
            'total_features': int(df_imputed.shape[1]),
            'missing_values_count': int(df_imputed.isnull().sum().sum()),
            'missing_rate': float(df_imputed.isnull().sum().sum() / (df_imputed.shape[0] * df_imputed.shape[1]))
        }

        # 确保所有数据都是可序列化的
        imputation_metrics = convert_to_serializable(imputation_metrics)

        # 构建返回结果
        result = {
            'status': 'success',
            'message': '随机森林填充完成',
            'imputed_data_path': output_csv_path,
            'original_stats': original_stats,
            'final_stats': final_stats,
            'imputation_metrics': imputation_metrics,
            'columns_imputed': list(imputation_metrics.keys()) if imputation_metrics else [],
            'imputation_method': 'Random Forest'
        }

        # 确保结果完全可序列化
        result = convert_to_serializable(result)

        return result

    except Exception as e:
        error_result = {
            'status': 'error',
            'message': f'随机森林填充失败: {str(e)}'
        }
        return error_result

if __name__ == "__main__":
    if len(sys.argv) < 2:
        result = {
            'status': 'error',
            'message': '缺少CSV文件路径参数'
        }
        print(json.dumps(result, ensure_ascii=False))
        sys.exit(1)

    csv_file_path = sys.argv[1]
    result = random_forest_imputation(csv_file_path)

    # 使用ensure_ascii=False和默认的indent来避免编码问题
    try:
        print(json.dumps(result, ensure_ascii=False))
    except Exception as e:
        # 如果仍有序列化问题，输出基本错误信息
        error_result = {
            'status': 'error',
            'message': f'结果序列化失败: {str(e)}'
        }
        print(json.dumps(error_result))
